package com.wsjzzcbq.filter;

import com.alibaba.fastjson2.JSONObject;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.factory.RequestRateLimiterGatewayFilterFactory;
import org.springframework.cloud.gateway.filter.ratelimit.KeyResolver;
import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter;
import org.springframework.cloud.gateway.route.Route;
import org.springframework.cloud.gateway.support.HttpStatusHolder;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;
import java.util.Iterator;
import java.util.Map;

/**
 * NewRequestRateLimiterGatewayFilterFactory
 *
 * @author wsjz
 * @date 2023/09/17
 */
@Component
public class NewRequestRateLimiterGatewayFilterFactory extends RequestRateLimiterGatewayFilterFactory {

    private final RateLimiter defaultRateLimiter;

    private final KeyResolver defaultKeyResolver;

    private boolean denyEmptyKey = true;

    private String emptyKeyStatusCode = HttpStatus.FORBIDDEN.name();

    public NewRequestRateLimiterGatewayFilterFactory(RateLimiter defaultRateLimiter, KeyResolver defaultKeyResolver) {
        super(defaultRateLimiter, defaultKeyResolver);
        this.defaultRateLimiter = defaultRateLimiter;
        this.defaultKeyResolver = defaultKeyResolver;
    }

    @Override
    public GatewayFilter apply(Config config) {
        System.out.println("过滤限流");
        KeyResolver resolver = (KeyResolver)this.getOrDefault(config.getKeyResolver(), this.defaultKeyResolver);
        RateLimiter<Object> limiter = (RateLimiter)this.getOrDefault(config.getRateLimiter(), this.defaultRateLimiter);
        boolean denyEmpty = (Boolean)this.getOrDefault(config.getDenyEmptyKey(), this.denyEmptyKey);
        HttpStatusHolder emptyKeyStatus = HttpStatusHolder.parse((String)this.getOrDefault(config.getEmptyKeyStatus(), this.emptyKeyStatusCode));
        return (exchange, chain) -> {
            return resolver.resolve(exchange).defaultIfEmpty("____EMPTY_KEY__").flatMap((key) -> {
                if ("____EMPTY_KEY__".equals(key)) {
                    if (denyEmpty) {
                        ServerWebExchangeUtils.setResponseStatus(exchange, emptyKeyStatus);
                        return exchange.getResponse().setComplete();
                    } else {
                        return chain.filter(exchange);
                    }
                } else {
                    String routeId = config.getRouteId();
                    if (routeId == null) {
                        Route route = (Route)exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_ROUTE_ATTR);
                        routeId = route.getId();
                    }

                    return limiter.isAllowed(routeId, key).flatMap((response) -> {
                        Iterator var4 = response.getHeaders().entrySet().iterator();

                        while(var4.hasNext()) {
                            Map.Entry<String, String> header = (Map.Entry)var4.next();
                            exchange.getResponse().getHeaders().add((String)header.getKey(), (String)header.getValue());
                        }

                        if (response.isAllowed()) {
                            return chain.filter(exchange);
                        } else {
                            //ServerWebExchangeUtils.setResponseStatus(exchange, config.getStatusCode());
                            ServerHttpResponse httpResponse = exchange.getResponse();
                            httpResponse.getHeaders().set("Content-Type", "application/json");
                            JSONObject json = new JSONObject();
                            json.put("code", 0);
                            json.put("msg", "当前请求人数较多，请稍后再访问");
                            DataBuffer dataBuffer = httpResponse.bufferFactory().wrap(json.toJSONString().getBytes(StandardCharsets.UTF_8));
                            return httpResponse.writeWith(Mono.just(dataBuffer));
                            //return exchange.getResponse().setComplete();
                        }
                    });
                }
            });
        };
    }

    private <T> T getOrDefault(T configValue, T defaultValue) {
        return configValue != null ? configValue : defaultValue;
    }
}
